/*
 * @Author: wuqingchun
 * @Date: 2024-05-23 15:11:18
 * @Description: 协程调度器
 * @LastEditors: wuqingchun
 * @LastEditTime: 2024-05-24 09:58:29
 */
#ifndef __CO_SCHEDULER_H__
#define __CO_SCHEDULER_H__

#include <unistd.h>
#include <sys/epoll.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>


#include "coroutine.h"
#include "co_connection.h"


typedef void (*CO_FUNC)(void* pScheduler, uint32_t uiCoId);

class CoScheduler
{
private:
    static constexpr size_t COROUTINE_COUNT = 1024;
    static constexpr size_t EVENT_MAX_COUNT = 1024;
    static constexpr int WAIT_TIMEOUT_MS = 10;   // epoll_wait 等待超时时间（单位为微秒）
public:
    CoScheduler()
    {

    }

    ~CoScheduler()
    {

    }
    

    int Init()
    {
        m_uiACoCtxId = CoAdd(AcceptCoroutine);
        m_uiRCoCtxId = CoAdd(ReadCoroutine);
        m_uiWCoCtxId = CoAdd(WriteCoroutine);
        if (m_uiACoCtxId == CoContext::INVALID_COROUTINE_ID ||
            m_uiRCoCtxId == CoContext::INVALID_COROUTINE_ID ||
            m_uiWCoCtxId == CoContext::INVALID_COROUTINE_ID)
        {
            return -1;
        }

        m_iEpollFd = epoll_create(1);
        if (m_iEpollFd == -1)
        {
            return -1;
        }

        return 0;
    }

    int AddConnection(Connection& refConn, uint32_t uiEvents)
    {
        return __ConnectionCtl(EPOLL_CTL_ADD, refConn, uiEvents);
    }

    int ModifyConnection(Connection& refConn, uint32_t uiEvents)
    {
        return __ConnectionCtl(EPOLL_CTL_MOD, refConn, uiEvents);
    }

    int DeleteConnection(Connection& refConn)
    {
        return __ConnectionCtl(EPOLL_CTL_DEL, refConn, 0);
    }

    int Schedule()
    {
        if (m_iEpollFd == -1)
        {
            return -1;
        }

        int iFdNum = epoll_wait(m_iEpollFd, m_astEvents, EVENT_MAX_COUNT, WAIT_TIMEOUT_MS);
        if (iFdNum == -1)
        {
            if (errno == EINTR)
            {
                return 0;
            }

            printf("epoll_wait error, %s\n", strerror(errno));

            return -1;
        }

        struct epoll_event* pstEvent;

        for (int i = 0; i < iFdNum; i++)
        {
            Connection* pobjConn = (Connection*)m_astEvents[i].data.ptr;
            uint32_t uiEvents = 0;
            printf("wait %d, %x\n", iFdNum, uiEvents);

            if ((uiEvents | EPOLLIN) == EPOLLIN)
            {
                if (pobjConn->IsListener())
                {
                    printf("Enter accept Connection %d\n", pobjConn->GetSock());
                    CoResume(m_uiACoCtxId, pobjConn);
                    printf("Exit accept Connection %d\n", pobjConn->GetSock());
                }
                else
                {
                    printf("Enter read Connection %d\n", pobjConn->GetSock());
                    CoResume(m_uiRCoCtxId, pobjConn);
                    printf("Exit read Connection %d\n", pobjConn->GetSock());
                }
            }

            if ((uiEvents | EPOLLOUT) == EPOLLOUT)
            {
                printf("Enter write Connection %d\n", pobjConn->GetSock());
                CoResume(m_uiWCoCtxId, pobjConn);
                printf("Exit write Connection %d\n", pobjConn->GetSock());
            }

            if ((uiEvents | EPOLLHUP | EPOLLERR) != 0)
            {
                printf("Connection close, %d\n", pobjConn->GetSock());
                pobjConn->Close();
                DeleteConnection(*pobjConn);
            }
        }

        return 0;
    }

    uint32_t CoAdd(CO_FUNC pfFunc)
    {
        if (m_uiSubCoCount == COROUTINE_COUNT)
        {
            return CoContext::INVALID_COROUTINE_ID;
        }

        CoContext* pstCo = new (std::nothrow) CoContext;
        if (pstCo == nullptr)
        {
            return CoContext::INVALID_COROUTINE_ID;
        }
    
        int iErr = getcontext(&pstCo->stUCtx);
        if (iErr == -1)
        {
            printf("getcontext error, %s\n", strerror(errno));
            delete pstCo;
            return CoContext::INVALID_COROUTINE_ID;
        }

        pstCo->stUCtx.uc_stack.ss_sp = pstCo->acStack;
        pstCo->stUCtx.uc_stack.ss_size = sizeof(pstCo->acStack);  
        pstCo->stUCtx.uc_link = &m_stMainUCtx; // 设置返回时切换到主流程
        pstCo->uiCtxId = m_uiNextID;
        m_arrSubCo[m_uiNextID] = pstCo;
        makecontext(&pstCo->stUCtx, (void (*)())pfFunc, 2, this, m_uiNextID);

        m_uiSubCoCount++;
        return m_uiNextID++;
    }

    void CoDelete(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return;
        }

        CoContext* pstCo = m_arrSubCo[uiCoId];
        if (pstCo == nullptr)
        {
            return;
        }

        delete pstCo;
        m_arrSubCo[uiCoId] = nullptr;
        m_uiSubCoCount--;

        if (m_uiNextID > uiCoId)
        {
            m_uiNextID = uiCoId;
        }
    }

    int CoResume(uint32_t uiCoId, void* pUserData)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return -1;
        }

        return __Resume(*m_arrSubCo[uiCoId], pUserData);
    }

    int CoYeild(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return -1;
        }
        
        return __Yield(*m_arrSubCo[uiCoId]);
    }

    void* GetCoUserData(uint32_t uiCoId)
    {
        if (uiCoId == CoContext::INVALID_COROUTINE_ID)
        {
            return nullptr;
        }

        return m_arrSubCo[uiCoId]->pUserData;
    }

private:
    int __ConnectionCtl(int iOperType, Connection& refConn, uint32_t uiEvents)
    {
        struct epoll_event stEvent;

        if (m_iEpollFd == -1)
        {
            return -1;
        }

        int iFd = refConn.GetSock();

        if (iFd == -1)
        {
            return -1;
        }

        stEvent.data.ptr = (void*)&refConn;
        stEvent.events = uiEvents;
        return epoll_ctl(m_iEpollFd, iOperType, iFd, &stEvent);
    }

    int __Yield(CoContext& refCoCtx)
    {
        if (refCoCtx.enState != CoState::CO_STATE_RUNNING)
        {
            return 0;
        }

        int iErr = swapcontext(&refCoCtx.stUCtx, &m_stMainUCtx);          // 切换到主协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }

    int __Resume(CoContext& refCoCtx, void* pUserData)
    {
        if (refCoCtx.enState == CoState::CO_STATE_RUNNING)
        {
            return 0;
        }
        
        refCoCtx.pUserData = pUserData;
        int iErr = swapcontext(&m_stMainUCtx, &refCoCtx.stUCtx);          // 切换到主协程
        if (iErr == -1)
        {
            printf("swapcontext error, %s\n", strerror(errno));
        }

        return iErr;
    }
private:
    ucontext_t m_stMainUCtx;                                       // 主协程
    std::array<CoContext*, COROUTINE_COUNT> m_arrSubCo;     // 子协程指针数组
    uint32_t m_uiSubCoCount = 0;
    uint32_t m_uiNextID = 0;
private:
    int m_iEpollFd = -1;
    struct epoll_event m_astEvents[EVENT_MAX_COUNT];
    uint32_t m_uiACoCtxId = CoContext::INVALID_COROUTINE_ID;
    uint32_t m_uiRCoCtxId = CoContext::INVALID_COROUTINE_ID;
    uint32_t m_uiWCoCtxId = CoContext::INVALID_COROUTINE_ID;
};


#endif //__COROUTINE_SCHEDULER_H__